{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tiny NNX\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/cgarciae/nnx/blob/main/docs/tiny_nnx.ipynb)\n",
    "\n",
    "A pedagogical implementation of NNX's core APIs.\n",
    "\n",
    "## Core APIs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dataclasses\n",
    "import hashlib\n",
    "import typing as tp\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import random\n",
    "\n",
    "A = tp.TypeVar(\"A\")\n",
    "M = tp.TypeVar(\"M\", bound=\"Module\")\n",
    "Sharding = tp.Tuple[tp.Optional[str], ...]\n",
    "Array = jax.Array\n",
    "\n",
    "\n",
    "class Variable(tp.Generic[A]):\n",
    "\n",
    "  def __init__(\n",
    "      self,\n",
    "      value: A,\n",
    "      *,\n",
    "      sharding: tp.Optional[Sharding] = None,\n",
    "  ):\n",
    "    self.value = value\n",
    "    self.sharding = sharding\n",
    "\n",
    "  def __repr__(self) -> str:\n",
    "    return (\n",
    "        f\"{type(self).__name__}(value={self.value}, sharding={self.sharding})\"\n",
    "    )\n",
    "\n",
    "  def __init_subclass__(cls):\n",
    "    super().__init_subclass__()\n",
    "    jax.tree_util.register_pytree_node(\n",
    "        cls,\n",
    "        lambda x: ((x.value,), (x.sharding,)),\n",
    "        lambda metadata, value: cls(value[0], sharding=metadata[0]),\n",
    "    )\n",
    "\n",
    "\n",
    "class State(dict[str, Variable[tp.Any]]):\n",
    "\n",
    "  def extract(self, variable_type: tp.Type[Variable]) -> \"State\":\n",
    "    return State(\n",
    "        {\n",
    "            path: variable\n",
    "            for path, variable in self.items()\n",
    "            if isinstance(variable, variable_type)\n",
    "        }\n",
    "    )\n",
    "\n",
    "  def __repr__(self) -> str:\n",
    "    elems = \",\\n  \".join(\n",
    "        f\"'{path}': {variable}\".replace(\"\\n\", \"\\n    \")\n",
    "        for path, variable in self.items()\n",
    "    )\n",
    "    return f\"State({{\\n  {elems}\\n}})\"\n",
    "\n",
    "\n",
    "jax.tree_util.register_pytree_node(\n",
    "    State,\n",
    "    # in reality, values and paths should be sorted by path\n",
    "    lambda x: (tuple(x.values()), tuple(x.keys())),\n",
    "    lambda paths, values: State(dict(zip(paths, values))),\n",
    ")\n",
    "\n",
    "\n",
    "@dataclasses.dataclass\n",
    "class GraphDef(tp.Generic[M]):\n",
    "  type: tp.Type[M]\n",
    "  index: int\n",
    "  submodules: dict[str, tp.Union[\"GraphDef[Module]\", int]]\n",
    "  static_fields: dict[str, tp.Any]\n",
    "\n",
    "  def merge(self, state: State) -> M:\n",
    "    module = GraphDef._build_module_recursive(self, {})\n",
    "    module.update(state)\n",
    "    return module\n",
    "\n",
    "  @staticmethod\n",
    "  def _build_module_recursive(\n",
    "      graphdef: tp.Union[\"GraphDef[M]\", int],\n",
    "      index_to_module: dict[int, \"Module\"],\n",
    "  ) -> M:\n",
    "    if isinstance(graphdef, int):\n",
    "      return index_to_module[graphdef] # type: ignore\n",
    "\n",
    "    assert graphdef.index not in index_to_module\n",
    "\n",
    "    # add a dummy module to the index to avoid infinite recursion\n",
    "    module = object.__new__(graphdef.type)\n",
    "    index_to_module[graphdef.index] = module\n",
    "\n",
    "    submodules = {\n",
    "        name: GraphDef._build_module_recursive(submodule, index_to_module)\n",
    "        for name, submodule in graphdef.submodules.items()\n",
    "    }\n",
    "    vars(module).update(graphdef.static_fields)\n",
    "    vars(module).update(submodules)\n",
    "    return module\n",
    "\n",
    "  def apply(\n",
    "      self, state: State\n",
    "  ) -> tp.Callable[..., tuple[tp.Any, tuple[State, \"GraphDef[M]\"]]]:\n",
    "    def _apply(*args, **kwargs):\n",
    "      module = self.merge(state)\n",
    "      out = module(*args, **kwargs)  # type: ignore\n",
    "      return out, module.split()\n",
    "\n",
    "    return _apply\n",
    "\n",
    "\n",
    "class Module:\n",
    "\n",
    "  def split(self: M) -> tp.Tuple[State, GraphDef[M]]:\n",
    "    state = State()\n",
    "    graphdef = Module._partition_recursive(\n",
    "        module=self, module_id_to_index={}, path_parts=(), state=state\n",
    "    )\n",
    "    assert isinstance(graphdef, GraphDef)\n",
    "    return state, graphdef\n",
    "\n",
    "  @staticmethod\n",
    "  def _partition_recursive(\n",
    "      module: M,\n",
    "      module_id_to_index: dict[int, int],\n",
    "      path_parts: tp.Tuple[str, ...],\n",
    "      state: State,\n",
    "  ) -> tp.Union[GraphDef[M], int]:\n",
    "    if id(module) in module_id_to_index:\n",
    "      return module_id_to_index[id(module)]\n",
    "\n",
    "    index = len(module_id_to_index)\n",
    "    module_id_to_index[id(module)] = index\n",
    "\n",
    "    submodules = {}\n",
    "    static_fields = {}\n",
    "\n",
    "    # iterate fields sorted by name to ensure deterministic order\n",
    "    for name, value in sorted(vars(module).items(), key=lambda x: x[0]):\n",
    "      value_path = (*path_parts, name)\n",
    "      # if value is a Module, recurse\n",
    "      if isinstance(value, Module):\n",
    "        submoduledef = Module._partition_recursive(\n",
    "            value, module_id_to_index, value_path, state\n",
    "        )\n",
    "        submodules[name] = submoduledef\n",
    "      # if value is a Variable, add to state\n",
    "      elif isinstance(value, Variable):\n",
    "        state[\"/\".join(value_path)] = value\n",
    "      else:  # otherwise, add to graphdef fields\n",
    "        static_fields[name] = value\n",
    "\n",
    "    return GraphDef(\n",
    "        type=type(module),\n",
    "        index=index,\n",
    "        submodules=submodules,\n",
    "        static_fields=static_fields,\n",
    "    )\n",
    "\n",
    "  def update(self, state: State) -> None:\n",
    "    for path, value in state.items():\n",
    "      path_parts = path.split(\"/\")\n",
    "      Module._set_value_at_path(self, path_parts, value)\n",
    "\n",
    "  @staticmethod\n",
    "  def _set_value_at_path(\n",
    "      module: \"Module\", path_parts: tp.Sequence[str], value: Variable[tp.Any]\n",
    "  ) -> None:\n",
    "    if len(path_parts) == 1:\n",
    "      setattr(module, path_parts[0], value)\n",
    "    else:\n",
    "      Module._set_value_at_path(\n",
    "          getattr(module, path_parts[0]), path_parts[1:], value\n",
    "      )\n",
    "\n",
    "\n",
    "@dataclasses.dataclass\n",
    "class Rngs:\n",
    "  key: jax.Array\n",
    "  count: int = 0\n",
    "  count_path: tuple[int, ...] = ()\n",
    "\n",
    "  def fork(self) -> \"Rngs\":\n",
    "    \"\"\"Forks the context, guaranteeing that all the random numbers generated\n",
    "    will be different from the ones generated in the original context. Fork is\n",
    "    used to create a new Rngs that can be passed to a JAX transform\"\"\"\n",
    "    count_path = self.count_path + (self.count,)\n",
    "    self.count += 1\n",
    "    return Rngs(self.key, count_path=count_path)\n",
    "\n",
    "  def make_rng(self) -> jax.Array:\n",
    "    fold_data = self._stable_hash(self.count_path + (self.count,))\n",
    "    self.count += 1\n",
    "    return random.fold_in(self.key, fold_data)  # type: ignore\n",
    "\n",
    "  @staticmethod\n",
    "  def _stable_hash(data: tuple[int, ...]) -> int:\n",
    "    hash_str = \" \".join(str(x) for x in data)\n",
    "    _hash = hashlib.blake2s(hash_str.encode())\n",
    "    hash_bytes = _hash.digest()\n",
    "    # uint32 is represented as 4 bytes in big endian\n",
    "    return int.from_bytes(hash_bytes[:4], byteorder=\"big\")\n",
    "\n",
    "\n",
    "# in the real NNX Rngs is not a pytree, instead\n",
    "# it has a split/merge API similar to Module\n",
    "# but for simplicity we use a pytree here\n",
    "jax.tree_util.register_pytree_node(\n",
    "    Rngs,\n",
    "    lambda x: ((x.key,), (x.count, x.count_path)),\n",
    "    lambda metadata, value: Rngs(value[0], *metadata),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Param(Variable[A]):\n",
    "  pass\n",
    "\n",
    "\n",
    "class BatchStat(Variable[A]):\n",
    "  pass\n",
    "\n",
    "\n",
    "class Linear(Module):\n",
    "\n",
    "  def __init__(self, din: int, dout: int, *, rngs: Rngs):\n",
    "    self.din = din\n",
    "    self.dout = dout\n",
    "    key = rngs.make_rng()\n",
    "    self.w = Param(random.uniform(key, (din, dout)))\n",
    "    self.b = Param(jnp.zeros((dout,)))\n",
    "\n",
    "  def __call__(self, x: jax.Array) -> jax.Array:\n",
    "    return x @ self.w.value + self.b.value\n",
    "\n",
    "\n",
    "class BatchNorm(Module):\n",
    "\n",
    "  def __init__(self, din: int, mu: float = 0.95):\n",
    "    self.mu = mu\n",
    "    self.scale = Param(jax.numpy.ones((din,)))\n",
    "    self.bias = Param(jax.numpy.zeros((din,)))\n",
    "    self.mean = BatchStat(jax.numpy.zeros((din,)))\n",
    "    self.var = BatchStat(jax.numpy.ones((din,)))\n",
    "\n",
    "  def __call__(self, x, train: bool) -> jax.Array:\n",
    "    if train:\n",
    "      axis = tuple(range(x.ndim - 1))\n",
    "      mean = jax.numpy.mean(x, axis=axis)\n",
    "      var = jax.numpy.var(x, axis=axis)\n",
    "      # ema update\n",
    "      self.mean.value = self.mu * self.mean.value + (1 - self.mu) * mean\n",
    "      self.var.value = self.mu * self.var.value + (1 - self.mu) * var\n",
    "    else:\n",
    "      mean, var = self.mean.value, self.var.value\n",
    "\n",
    "    scale, bias = self.scale.value, self.bias.value\n",
    "    x = (x - mean) / jax.numpy.sqrt(var + 1e-5) * scale + bias\n",
    "    return x\n",
    "\n",
    "\n",
    "class Dropout(Module):\n",
    "\n",
    "  def __init__(self, rate: float):\n",
    "    self.rate = rate\n",
    "\n",
    "  def __call__(self, x: jax.Array, *, train: bool, rngs: Rngs) -> jax.Array:\n",
    "    if train:\n",
    "      mask = random.bernoulli(rngs.make_rng(), (1 - self.rate), x.shape)\n",
    "      x = x * mask / (1 - self.rate)\n",
    "    return x"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scan Over Layers Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(Module):\n",
    "\n",
    "  def __init__(self, din: int, dout: int, *, rngs: Rngs):\n",
    "    self.linear = Linear(din, dout, rngs=rngs)\n",
    "    self.bn = BatchNorm(dout)\n",
    "    self.dropout = Dropout(0.1)\n",
    "\n",
    "  def __call__(self, x: jax.Array, *, train: bool, rngs: Rngs) -> jax.Array:\n",
    "    x = self.linear(x)\n",
    "    x = self.bn(x, train=train)\n",
    "    x = jax.nn.gelu(x)\n",
    "    x = self.dropout(x, train=train, rngs=rngs)\n",
    "    return x\n",
    "\n",
    "\n",
    "class ScanMLP(Module):\n",
    "\n",
    "  def __init__(self, hidden_size: int, n_layers: int, *, rngs: Rngs):\n",
    "    self.n_layers = n_layers\n",
    "\n",
    "    # lift init\n",
    "    key = random.split(rngs.make_rng(), n_layers - 1)\n",
    "    graphdef: GraphDef[Block] = None  # type: ignore\n",
    "\n",
    "    def init_fn(key):\n",
    "      nonlocal graphdef\n",
    "      state, graphdef = Block(\n",
    "          hidden_size, hidden_size, rngs=Rngs(key)\n",
    "      ).split()\n",
    "      return state\n",
    "\n",
    "    state = jax.vmap(init_fn)(key)\n",
    "    self.layers = graphdef.merge(state)\n",
    "    self.linear = Linear(hidden_size, hidden_size, rngs=rngs)\n",
    "\n",
    "  def __call__(self, x: jax.Array, *, train: bool, rngs: Rngs) -> jax.Array:\n",
    "    # lift call\n",
    "    key: jax.Array = random.split(rngs.make_rng(), self.n_layers - 1)  # type: ignore\n",
    "    state, graphdef = self.layers.split()\n",
    "\n",
    "    def scan_fn(x, inputs: tuple[jax.Array, State]):\n",
    "      key, state = inputs\n",
    "      x, (state, _) = graphdef.apply(state)(x, train=train, rngs=Rngs(key))\n",
    "      return x, state\n",
    "\n",
    "    x, state = jax.lax.scan(scan_fn, x, (key, state))\n",
    "    self.layers.update(state)\n",
    "    x = self.linear(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "state = State({\n",
      "  'layers/bn/bias': Param(value=(4, 10), sharding=None),\n",
      "  'layers/bn/mean': BatchStat(value=(4, 10), sharding=None),\n",
      "  'layers/bn/scale': Param(value=(4, 10), sharding=None),\n",
      "  'layers/bn/var': BatchStat(value=(4, 10), sharding=None),\n",
      "  'layers/linear/b': Param(value=(4, 10), sharding=None),\n",
      "  'layers/linear/w': Param(value=(4, 10, 10), sharding=None),\n",
      "  'linear/b': Param(value=(10,), sharding=None),\n",
      "  'linear/w': Param(value=(10, 10), sharding=None)\n",
      "})\n",
      "graphdef = GraphDef(type=<class '__main__.ScanMLP'>, index=0, submodules={'layers': GraphDef(type=<class '__main__.Block'>, index=1, submodules={'bn': GraphDef(type=<class '__main__.BatchNorm'>, index=2, submodules={}, static_fields={'mu': 0.95}), 'dropout': GraphDef(type=<class '__main__.Dropout'>, index=3, submodules={}, static_fields={'rate': 0.1}), 'linear': GraphDef(type=<class '__main__.Linear'>, index=4, submodules={}, static_fields={'din': 10, 'dout': 10})}, static_fields={}), 'linear': GraphDef(type=<class '__main__.Linear'>, index=5, submodules={}, static_fields={'din': 10, 'dout': 10})}, static_fields={'n_layers': 5})\n"
     ]
    }
   ],
   "source": [
    "module = ScanMLP(hidden_size=10, n_layers=5, rngs=Rngs(random.key(0)))\n",
    "x = jax.random.normal(random.key(0), (2, 10))\n",
    "y = module(x, train=True, rngs=Rngs(random.key(1)))\n",
    "\n",
    "state, graphdef = module.split()\n",
    "print(\"state =\", jax.tree.map(jnp.shape, state))\n",
    "print(\"graphdef =\", graphdef)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filtering State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "params = State({\n",
      "  'layers/bn/bias': Param(value=(4, 10), sharding=None),\n",
      "  'layers/bn/scale': Param(value=(4, 10), sharding=None),\n",
      "  'layers/linear/b': Param(value=(4, 10), sharding=None),\n",
      "  'layers/linear/w': Param(value=(4, 10, 10), sharding=None),\n",
      "  'linear/b': Param(value=(10,), sharding=None),\n",
      "  'linear/w': Param(value=(10, 10), sharding=None)\n",
      "})\n",
      "batch_stats = State({\n",
      "  'layers/bn/mean': BatchStat(value=(4, 10), sharding=None),\n",
      "  'layers/bn/var': BatchStat(value=(4, 10), sharding=None)\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "# split\n",
    "params = state.extract(Param)\n",
    "batch_stats = state.extract(BatchStat)\n",
    "# merge\n",
    "state = State({**params, **batch_stats})\n",
    "\n",
    "print(\"params =\", jax.tree.map(jnp.shape, params))\n",
    "print(\"batch_stats =\", jax.tree.map(jnp.shape, batch_stats))"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
